#include <iostream>
#include <sstream>
#include <fstream>
#include <QApplication>
#include <QThread>
#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>

#include "base/pose_io.hh"
#include "base/global_param.hh"
#include "base/frame_grabber.hh"
#include "base/util.hh"
#include "mbt/object3d.hh"
#include "mbt/tracker_rbot.hh"
#include "mbt/tracker_sle.hh"

void ConvertMask(const cv::Mat& src_mask, cv::Mat& mask, uchar oid) {
	mask = cv::Mat(src_mask.size(), CV_8UC1, cv::Scalar(0));
	uchar depth = src_mask.type() & CV_MAT_DEPTH_MASK;

	if (CV_8U == depth && oid > 0) {
		for (int r = 0; r < src_mask.rows; ++r)
		for (int c = 0; c < src_mask.cols; ++c) {
			if (oid == src_mask.at<uchar>(r,c))
				mask.at<uchar>(r,c) = 255;
		}
	} else 
	if (CV_32F == depth) {
		for (int r = 0; r < src_mask.rows; ++r)
		for (int c = 0; c < src_mask.cols; ++c) {
			if (src_mask.at<float>(r,c))
				mask.at<uchar>(r,c) = 255;
		}
	}	else {
		LOG(ERROR) << "WRONG IMAGE TYPE";
	}
}

class RecBase {
public:
	virtual void SaveFrame(std::vector<Model*>& objects, int fid, cv::Mat& frame, View* view) = 0;
};

class PlayRec : public RecBase {
public:
	PlayRec() {
		tk::GlobalParam* gp = tk::GlobalParam::Instance();
	}

	virtual void SaveFrame(std::vector<Model*>& objects, int fid, cv::Mat& frame, View* view) override {
		cv::Mat result = View::Instance()->DrawMeshOverlay(std::vector<Model*>(objects.begin(), objects.end()), frame);
		std::stringstream sstr;
		sstr << "#" << std::setw(4) << std::setfill('0') << fid;
		cv::putText(result, sstr.str(), cv::Point(5, 20), cv::FONT_HERSHEY_COMPLEX_SMALL, 1.0, cv::Scalar(255, 255, 255), 1);
		cv::imshow("result", result);
		cv::waitKey();
	}
};

class MaskRec: public RecBase {
public:
	MaskRec() {}

	virtual void SaveFrame(std::vector<Model*>& objects, int fid, cv::Mat& frame, View* view) override {
		tk::GlobalParam* gp = tk::GlobalParam::Instance();

		view->RenderSilhouette(std::vector<Model*>(objects.begin(), objects.end()), GL_FILL);
		cv::Mat masks = view->DownloadFrame(View::MASK);
		cv::Mat mask;
		ConvertMask(masks, mask, 1);
		//cv::imshow("mask", mask);
		
		std::stringstream sstr1;
		sstr1 << gp->bench_case + "/mask_" << std::setw(4) << std::setfill('0') << fid << ".png";
		cv::imwrite(sstr1.str(), mask);
	}
};

class MeshRec: public RecBase {
public:
	MeshRec() {}

	virtual void SaveFrame(std::vector<Model*>& objects, int fid, cv::Mat& frame, View* view) override {
		tk::GlobalParam* gp = tk::GlobalParam::Instance();

		cv::Mat result = View::Instance()->DrawMeshOverlay(std::vector<Model*>(objects.begin(), objects.end()), frame);

		std::stringstream sstr;
		sstr << "#" << std::setw(4) << std::setfill('0') << fid;
		cv::putText(result, sstr.str(), cv::Point(5, 20), cv::FONT_HERSHEY_COMPLEX_SMALL, 1.0, cv::Scalar(255, 255, 255), 1);

		std::stringstream str;
		//str << gp->bench_case + "/mesh_" << std::setw(4) << std::setfill('0') << fid << ".png";
		str << gp->bench_case + "/cam_" << std::setw(4) << std::setfill('0') << fid << ".png";
		cv::imwrite(str.str(), result);
	}
};

class ShadeRec: public RecBase {
public:
	ShadeRec() {}

	virtual void SaveFrame(std::vector<Model*>& objects, int fid, cv::Mat& frame, View* view) override {
		tk::GlobalParam* gp = tk::GlobalParam::Instance();

		cv::Mat result = View::Instance()->DrawResultOverlay(std::vector<Model*>(objects.begin(), objects.end()), frame);
		
		std::stringstream sstr;
		sstr << "#" << std::setw(4) << std::setfill('0') << fid;
		cv::putText(result, sstr.str(), cv::Point(5, 20), cv::FONT_HERSHEY_COMPLEX_SMALL, 1.0, cv::Scalar(255, 255, 255), 1);

		std::stringstream str;
		//str << gp->bench_case + "/shaded_" << std::setw(4) << std::setfill('0') << fid << ".png";
		str << gp->bench_case + "/cam_" << std::setw(4) << std::setfill('0') << fid << ".png";
		cv::imwrite(str.str(), result);
		//cv::imshow("shade", result);
		//cv::waitKey();
	}
};

class ContourRec: public RecBase {
public:
	ContourRec() {}

	virtual void SaveFrame(std::vector<Model*>& objects, int fid, cv::Mat& frame, View* view) override {
		tk::GlobalParam* gp = tk::GlobalParam::Instance();

		view->RenderSilhouette(std::vector<Model*>(objects.begin(), objects.end()), GL_FILL);

		cv::Mat depth_map = view->DownloadFrame(View::DEPTH);
		cv::Mat masks_map;
		if (objects.size() > 1) {
			masks_map = view->DownloadFrame(View::MASK);
		}	else {
			masks_map = depth_map;
		}

		for (int oid = 0; oid < objects.size(); oid++) {
			cv::Mat mask_map;
			ConvertMask(masks_map, mask_map, objects[oid]->getModelID());

			std::vector<std::vector<cv::Point> > contours;
			cv::findContours(mask_map, contours, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_NONE);

			cv::Vec3b color;
			if (0 == oid)
				color = cv::Vec3b(0, 255, 0);
			if (1 == oid)
				color = cv::Vec3b(0, 0, 255);

			for (auto contour : contours)
			for (auto pt : contour) {
				frame.at<cv::Vec3b>(pt) = color;
				//cv::circle(frame, pt, 2, color);
				//cv::circle(edge_c3, pt, 1, color);
			}
		}
		std::stringstream sstr;
		sstr << "#" << std::setw(3) << std::setfill('0') << fid;
		cv::putText(frame, sstr.str(), cv::Point(5, 20), cv::FONT_HERSHEY_COMPLEX_SMALL, 1.0, cv::Scalar(0, 255, 0), 1);

		std::stringstream str;
		str << gp->bench_case + "/contour_" << std::setw(4) << std::setfill('0') << fid << ".png";
		cv::imwrite(str.str(), frame);
	}

};

int MixResult(const std::string& config_file) {
	std::cout << std::endl << "#######################" << std::endl;
	tk::GlobalParam* gp = tk::GlobalParam::Instance();
	gp->ParseConfig(config_file);

	std::vector<int> fids;
	std::vector<std::vector<cv::Matx44f> > tk_poses;
	PoseReader* pose_reader = new PoseReaderRBOT;
	pose_reader->Read(gp->tk_pose_file, fids, tk_poses);

	//tk::FrameGrabber* grabber = tk::FrameGrabber::GetFrameGrabber("data/out/cam/out");
	tk::FrameGrabber* grabber = tk::FrameGrabber::GetFrameGrabber(gp->frames);
	CHECK(grabber) << "ERROR: check <frames> in config file.";
	gp->image_width = grabber->width;
	gp->image_height = grabber->height;
	cv::Matx33f K = cv::Matx33f(gp->fx, 0, gp->cx, 0, gp->fy, gp->cy, 0, 0, 1);
	cv::Matx14f distCoeffs = cv::Matx14f(0.0, 0.0, 0.0, 0.0);

	View* view = View::Instance();
	view->init(K, gp->image_width, gp->image_height, gp->zn, gp->zf, 4);

	std::vector<float> distances = { 200.0f, 400.0f, 600.0f };
	std::vector<Object3D*> objects;
	for (int i = 0; i < gp->model_file.size(); ++i) {
		objects.push_back(new Object3D(gp->model_file[i], cv::Matx44f(), 1.0, 0.55f, distances));
	}

	Tracker* poseEstimator = new SLETracker(K, objects);

	for (int i = 0; i < objects.size(); ++i) {
		objects[i]->initialize();
	}

	RecBase* rec;

	//rec = new MaskRec();
	rec = new ShadeRec();
	//rec = new MeshRec();
	//rec = new PlayRec();
	//rec = new ContourRec();

	cv::Mat frame;
	for (int i = 0; i < fids.size(); ++i) {
		grabber->GrabFrame(frame, fids[i]);

		for (int oid = 0; oid < objects.size(); ++oid) {
			objects[oid]->setPose(tk_poses[oid][i]);
		}

		rec->SaveFrame(std::vector<Model*>(objects.begin(), objects.end()), fids[i], frame, view);
	}

	std::cout << "#######################" << std::endl;

	View::Instance()->destroy();

	for (int i = 0; i < objects.size(); i++) {
		delete objects[i];
	}
	objects.clear();
}

int main(int argc, char* argv[]) {
	QApplication a(argc, argv);

	MixResult(argv[1]);
}